/**
 * 
 */
package sjsu.cs297.etm;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Formatter;
import java.util.Iterator;
import java.util.Locale;
import java.util.TreeSet;

import javax.swing.JTable.PrintMode;

import org.junit.Assert;

import cc.mallet.topics.ParallelTopicModel;
import cc.mallet.topics.TopicInferencer;
import cc.mallet.types.Alphabet;
import cc.mallet.types.FeatureSequence;
import cc.mallet.types.IDSorter;
import cc.mallet.types.Instance;
import cc.mallet.types.InstanceList;
import cc.mallet.types.LabelSequence;

/**
 * Main entry to run the application for email topic modeling.
 * 
 * @author Hiep.Hong
 * 
 */
public class ApplicationEntry {

	private InstanceList instances;
	private ParallelTopicModel model;

	public ApplicationEntry(String input, String output) {
		importData(input, output);
	}

	private void importData(String input, String output) {
		FileImporter importer = new FileImporter();

		printStatus("***Importing raw data from: " + input + "...");
		instances = importer.readDirectory(new File(input));

		printStatus("***Saving prepared data to: " + output + "...");
		instances.save(new File(output));
	}

	public void train(int numTopics, int numThreads, int numIterations) {
		try {
			printStatus("***Number of topics specified: " + numTopics);
			printStatus("***Number of threads to run training: " + numThreads);
			printStatus("***Begin training the data...");
			model = new ParallelTopicModel(numTopics);
			model.addInstances(instances);

			// Use two parallel samplers, which each look at one half the corpus
			// and combine statistics after every iteration.
			model.setNumThreads(numThreads);

			// Run the model for 50 iterations and stop (this is for testing
			// only, for real applications, use 1000 to 2000 iterations).
			model.setNumIterations(numIterations);
			model.estimate();

			printStatus("***Finish training");

			// Show the words and topics in the first instance

			// The data alphabet maps word IDs to strings
			Formatter out;
			Alphabet dataAlphabet = showWordsAndTopcis();

			showTopFiveWords(numTopics, dataAlphabet);

			// Get an array of sorted sets of word ID/count pairs
			ArrayList<TreeSet<IDSorter>> topicSortedWords = model
					.getSortedWords();
			// Create a new instance with high probability of topic 0
			StringBuilder topicZeroText = new StringBuilder();
			Iterator<IDSorter> iterator = topicSortedWords.get(0).iterator();

			int rank = 0;
			while (iterator.hasNext() && rank < 5) {
				IDSorter idCountPair = iterator.next();
				topicZeroText.append(dataAlphabet.lookupObject(idCountPair
						.getID()) + " ");
				rank++;
			}

			// Create a new instance named "test instance" with empty target and
			// source fields.
			// InstanceList testing = instances;

			// TopicInferencer inferencer = model.getInferencer();
			// double[] testProbabilities = inferencer.getSampledDistribution(
			// testing.get(0), 10, 1, 5);
			// System.out.println("0\t" + testProbabilities[0]);

		} catch (IOException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
	}

	private void showTopFiveWords(int numTopics, Alphabet dataAlphabet) {
		Formatter out;
		// Estimate the topic distribution of the first instance,
		// given the current Gibbs state.
		double[] topicDistribution = model.getTopicProbabilities(0);

		// Get an array of sorted sets of word ID/count pairs
		ArrayList<TreeSet<IDSorter>> topicSortedWords = model.getSortedWords();

		// Show top 5 words in topics with proportions for the first
		// document
		for (int topic = 0; topic < numTopics; topic++) {
			Iterator<IDSorter> iterator = topicSortedWords.get(topic)
					.iterator();

			out = new Formatter(new StringBuilder(), Locale.US);
			out.format("%d\t%.3f\t", topic, topicDistribution[topic]);
			int rank = 0;
			while (iterator.hasNext() && rank < 5) {
				IDSorter idCountPair = iterator.next();
				out.format("%s (%.0f) ",
						dataAlphabet.lookupObject(idCountPair.getID()),
						idCountPair.getWeight());
				rank++;
			}
			System.out.println(out);
		}
	}

	public Alphabet showWordsAndTopcis() {
		Alphabet dataAlphabet = instances.getDataAlphabet();

		FeatureSequence tokens = (FeatureSequence) model.getData().get(0).instance
				.getData();
		LabelSequence topics = model.getData().get(0).topicSequence;

		Formatter out = new Formatter(new StringBuilder(), Locale.US);
		for (int position = 0; position < tokens.getLength(); position++) {
			out.format("%s-%d ", dataAlphabet.lookupObject(tokens
					.getIndexAtPosition(position)), topics
					.getIndexAtPosition(position));
		}

		System.out.println(out);

		return dataAlphabet;
	}

	/**
	 * @param args
	 */
	public static void main(String[] args) {
		String input = args[0];
		String output = args[1];
		String numTopicsString = args[2];

		Assert.assertNotNull(input);
		Assert.assertNotNull(output);
		Assert.assertNotNull(numTopicsString);

		int numThreads = Integer.parseInt(args[3]);
		int numIterations = Integer.parseInt(args[4]);

		ApplicationEntry appEntry = new ApplicationEntry(input, output);
		appEntry.train(Integer.parseInt(numTopicsString), numThreads,
				numIterations);
	}

	private void printStatus(String message) {
		System.out.println(message);
	}

}
